import time
import os
import numpy as np
from itertools import chain
import torch
from tensorboardX import SummaryWriter

from utils.separated_buffer import SeparatedReplayBuffer
from utils.util import update_linear_schedule


def _t2n(x):
    """
    将PyTorch张量转换为NumPy数组
    
    参数:
        x: PyTorch张量
        
    返回:
        NumPy数组
    """
    return x.detach().cpu().numpy()


class Runner(object):
    """
    运行器基类，用于管理训练和评估过程
    """
    def __init__(self, config):
        """
        初始化运行器
        
        参数:
            config: 配置字典，包含所有参数和环境
        """
        self.all_args = config["all_args"]  # 所有参数
        self.envs = config["envs"]  # 训练环境
        self.eval_envs = config["eval_envs"]  # 评估环境
        self.device = config["device"]  # 计算设备
        self.num_agents = config["num_agents"]  # 智能体数量

        # 参数
        self.env_name = self.all_args.env_name  # 环境名称
        self.algorithm_name = self.all_args.algorithm_name  # 算法名称
        self.experiment_name = self.all_args.experiment_name  # 实验名称
        self.use_centralized_V = self.all_args.use_centralized_V  # 是否使用中心化价值函数
        self.use_obs_instead_of_state = self.all_args.use_obs_instead_of_state  # 是否使用观察代替状态
        self.num_env_steps = self.all_args.num_env_steps  # 环境步数
        self.episode_length = self.all_args.episode_length  # 回合长度
        self.n_rollout_threads = self.all_args.n_rollout_threads  # 并行环境数量
        self.n_eval_rollout_threads = self.all_args.n_eval_rollout_threads  # 评估并行环境数量
        self.use_linear_lr_decay = self.all_args.use_linear_lr_decay  # 是否使用线性学习率衰减
        self.hidden_size = self.all_args.hidden_size  # 隐藏层大小
        self.use_render = self.all_args.use_render  # 是否渲染
        self.recurrent_N = self.all_args.recurrent_N  # 循环层数

        # 间隔
        self.save_interval = self.all_args.save_interval  # 保存间隔
        self.use_eval = self.all_args.use_eval  # 是否使用评估
        self.eval_interval = self.all_args.eval_interval  # 评估间隔
        self.log_interval = self.all_args.log_interval  # 日志间隔

        # 目录
        self.model_dir = self.all_args.model_dir  # 模型目录

        if self.use_render:
            import imageio

            self.run_dir = config["run_dir"]  # 运行目录
            self.gif_dir = str(self.run_dir / "gifs")  # GIF目录
            if not os.path.exists(self.gif_dir):
                os.makedirs(self.gif_dir)  # 创建GIF目录
        else:
            # if self.use_wandb:
            #     self.save_dir = str(wandb.run.dir)
            # else:
            self.run_dir = config["run_dir"]  # 运行目录
            self.log_dir = str(self.run_dir / "logs")  # 日志目录
            if not os.path.exists(self.log_dir):
                os.makedirs(self.log_dir)  # 创建日志目录
            self.writter = SummaryWriter(self.log_dir)  # 创建TensorBoard写入器
            self.save_dir = str(self.run_dir / "models")  # 保存目录
            if not os.path.exists(self.save_dir):
                os.makedirs(self.save_dir)  # 创建保存目录

        from algorithms.algorithm.r_mappo import RMAPPO as TrainAlgo
        from algorithms.algorithm.rMAPPOPolicy import RMAPPOPolicy as Policy

        self.policy = []  # 策略列表
        for agent_id in range(self.num_agents):
            share_observation_space = (
                self.envs.share_observation_space[agent_id]
                if self.use_centralized_V
                else self.envs.observation_space[agent_id]
            )  # 共享观察空间
            # 策略网络
            po = Policy(
                self.all_args,
                self.envs.observation_space[agent_id],
                share_observation_space,
                self.envs.action_space[agent_id],
                device=self.device,
            )
            self.policy.append(po)  # 添加策略

        if self.model_dir is not None:
            self.restore()  # 恢复模型

        self.trainer = []  # 训练器列表
        self.buffer = []  # 缓冲区列表
        for agent_id in range(self.num_agents):
            # 算法
            tr = TrainAlgo(self.all_args, self.policy[agent_id], device=self.device)
            # 缓冲区
            share_observation_space = (
                self.envs.share_observation_space[agent_id]
                if self.use_centralized_V
                else self.envs.observation_space[agent_id]
            )  # 共享观察空间
            bu = SeparatedReplayBuffer(
                self.all_args,
                self.envs.observation_space[agent_id],
                share_observation_space,
                self.envs.action_space[agent_id],
            )
            self.buffer.append(bu)  # 添加缓冲区
            self.trainer.append(tr)  # 添加训练器

    def run(self):
        """
        运行训练过程
        """
        raise NotImplementedError

    def warmup(self):
        """
        预热阶段，收集初始数据
        """
        raise NotImplementedError

    def collect(self, step):
        """
        收集数据
        
        参数:
            step: 当前步数
        """
        raise NotImplementedError

    def insert(self, data):
        """
        插入数据到缓冲区
        
        参数:
            data: 要插入的数据
        """
        raise NotImplementedError

    @torch.no_grad()
    def compute(self):
        """
        计算回报
        """
        for agent_id in range(self.num_agents):
            self.trainer[agent_id].prep_rollout()  # 准备收集数据
            next_value = self.trainer[agent_id].policy.get_values(
                self.buffer[agent_id].share_obs[-1],
                self.buffer[agent_id].rnn_states_critic[-1],
                self.buffer[agent_id].masks[-1],
            )  # 获取下一个状态的价值
            next_value = _t2n(next_value)  # 转换为NumPy数组
            self.buffer[agent_id].compute_returns(next_value, self.trainer[agent_id].value_normalizer)  # 计算回报

    def train(self):
        """
        训练智能体
        
        返回:
            训练信息列表
        """
        train_infos = []
        for agent_id in range(self.num_agents):
            self.trainer[agent_id].prep_training()  # 准备训练
            train_info = self.trainer[agent_id].train(self.buffer[agent_id])  # 训练智能体
            train_infos.append(train_info)  # 添加训练信息
            self.buffer[agent_id].after_update()  # 更新后处理

        return train_infos

    def save(self):
        """
        保存模型
        """
        for agent_id in range(self.num_agents):
            policy_actor = self.trainer[agent_id].policy.actor  # 获取策略的演员网络
            torch.save(
                policy_actor.state_dict(),
                str(self.save_dir) + "/actor_agent" + str(agent_id) + ".pt",
            )  # 保存演员网络
            policy_critic = self.trainer[agent_id].policy.critic  # 获取策略的评论家网络
            torch.save(
                policy_critic.state_dict(),
                str(self.save_dir) + "/critic_agent" + str(agent_id) + ".pt",
            )  # 保存评论家网络

    def restore(self):
        """
        恢复模型
        """
        for agent_id in range(self.num_agents):
            policy_actor_state_dict = torch.load(str(self.model_dir) + "/actor_agent" + str(agent_id) + ".pt")  # 加载演员网络状态字典
            self.policy[agent_id].actor.load_state_dict(policy_actor_state_dict)  # 加载演员网络
            policy_critic_state_dict = torch.load(
                str(self.model_dir) + "/critic_agent" + str(agent_id) + ".pt"
            )  # 加载评论家网络状态字典
            self.policy[agent_id].critic.load_state_dict(policy_critic_state_dict)  # 加载评论家网络

    def log_train(self, train_infos, total_num_steps):
        """
        记录训练信息
        
        参数:
            train_infos: 训练信息列表
            total_num_steps: 总步数
        """
        for agent_id in range(self.num_agents):
            for k, v in train_infos[agent_id].items():
                agent_k = "agent%i/" % agent_id + k  # 智能体键
                # if self.use_wandb:
                #     pass
                # wandb.log({agent_k: v}, step=total_num_steps)
                # else:
                self.writter.add_scalars(agent_k, {agent_k: v}, total_num_steps)  # 添加标量

    def log_env(self, env_infos, total_num_steps):
        """
        记录环境信息
        
        参数:
            env_infos: 环境信息字典
            total_num_steps: 总步数
        """
        for k, v in env_infos.items():
            if len(v) > 0:
                # if self.use_wandb:
                #     wandb.log({k: np.mean(v)}, step=total_num_steps)
                # else:
                self.writter.add_scalars(k, {k: np.mean(v)}, total_num_steps)  # 添加标量
